import glm
import OpenGL.GL as gl
import numpy as np

from pygl.shader import Shader

class Program:
    def __init__(self, shaders: list[Shader]) -> None:
        self.program_id = gl.glCreateProgram()
        for shader in shaders:
            gl.glAttachShader(self.program_id, shader.shader_id)
        gl.glLinkProgram(self.program_id)
        status = gl.glGetProgramiv(self.program_id, gl.GL_LINK_STATUS)
        if status == gl.GL_FALSE:
            log = gl.glGetProgramInfoLog(self.program_id)
            print("Program {} Log: {}".format(self.program_id, log))   

    def use(self) -> None:
        gl.glUseProgram(self.program_id)

    def delete(self) -> None:
        if gl.glIsProgram(self.program_id):
            gl.glDeleteProgram(self.program_id)

    def setUniform1f(self, location:int, value: float) -> None:
        gl.glProgramUniform1f(self.program_id,location, value)

    def setUniform2f(self, location:int, v0: float, v1: float) -> None:
        gl.glProgramUniform2f(self.program_id,location, v0, v1)

    def setUniform3f(self, location:int, v0: float, 
                     v1: float, v2: float) -> None:
        gl.glProgramUniform3f(self.program_id,location, v0, v1, v2)

    def setUniform4f(self, location:int, v0: float, 
                     v1: float, v2: float, v3: float) -> None:
        gl.glProgramUniform4f(self.program_id,location, v0, v1, v2, v3)

    def setUniform1i(self, location:int, value: int) -> None:
        gl.glProgramUniform1i(self.program_id,location, value)

    def setUniform2fv(self, location:int, value:glm.vec2) -> None:
        gl.glProgramUniform2fv(self.program_id,location, 
                               1, glm.value_ptr(value))
    
    def setUniform3fv(self, location:int, value:glm.vec3) -> None:
        gl.glProgramUniform3fv(self.program_id,location, 
                               1, glm.value_ptr(value))

    def setUniform4fv(self, location:int, value:glm.vec4) -> None:
        gl.glProgramUniform4fv(self.program_id,location, 
                               1, glm.value_ptr(value))

    def setUniformMatrix4fv(self, location:int, mat: glm.mat4x4) -> None:
        gl.glProgramUniformMatrix4fv(self.program_id,location, 
                                     1, False, glm.value_ptr(mat))

    def setUniformMatrix3fv(self, location:int, mat: glm.mat3x3) -> None:
        gl.glProgramUniformMatrix4fv(self.program_id,location, 
                                     1, False, glm.value_ptr(mat))

    def setUniform1fv(self, location:int, count:int, value:list) -> None:
        gl.glProgramUniform1fv(self.program_id,location, 
                               count, value)

class ProgramVF(Program):
    def __init__(self, vertex_source_path:str, fragment_source_path:str) -> None:
        vertex_shader = Shader(gl.GL_VERTEX_SHADER, vertex_source_path)
        fragment_shader = Shader(gl.GL_FRAGMENT_SHADER, fragment_source_path)
        super().__init__([vertex_shader, fragment_shader])
        vertex_shader.delete()
        fragment_shader.delete()

class ProgramCompute(Program):
    def __init__(self, compute_source_path:str) -> None:
        compute_shader = Shader(gl.GL_COMPUTE_SHADER, compute_source_path)
        super().__init__([compute_shader])
        compute_shader.delete()
        self.local_size = np.zeros(3, dtype=np.int32)
        gl.glGetProgramiv(self.program_id,
                       gl.GL_COMPUTE_WORK_GROUP_SIZE, self.local_size)
        
    def dispatch(self, nx:int, ny:int, nz:int) -> None:
        gl.glDispatchCompute(nx, ny, nz)

class ProgramVFG(Program):
    def __init__(self, vertex_source_path:str, fragment_source_path:str, 
                 geometry_source_path:str) -> None:
        vertex_shader = Shader(gl.GL_VERTEX_SHADER, vertex_source_path)
        fragment_shader = Shader(gl.GL_FRAGMENT_SHADER, fragment_source_path)
        geometry_shader = Shader(gl.GL_GEOMETRY_SHADER, geometry_source_path)
        super().__init__([vertex_shader, fragment_shader, geometry_shader])
        vertex_shader.delete()
        fragment_shader.delete()
        geometry_shader.delete()

class ProgramVFT(Program):
    def __init__(self, vertex_source_path:str, fragment_source_path:str, 
                 tcs_source_path:str, tes_source_path:str) -> None:
        vertex_shader = Shader(gl.GL_VERTEX_SHADER, vertex_source_path)
        fragment_shader = Shader(gl.GL_FRAGMENT_SHADER, fragment_source_path)
        tcs_shader = Shader(gl.GL_TESS_CONTROL_SHADER, tcs_source_path)
        tes_shader = Shader(gl.GL_TESS_EVALUATION_SHADER, tes_source_path)
        super().__init__([vertex_shader, fragment_shader, tcs_shader, 
                          tes_shader])
        vertex_shader.delete()
        fragment_shader.delete()
        tcs_shader.delete()
        tes_shader.delete()

class ProgramVFTG(Program):
    def __init__(self, vertex_source_path:str, fragment_source_path:str, 
                 tcs_source_path:str, tes_source_path:str,
                 geometry_source_path:str) -> None:
        vertex_shader = Shader(gl.GL_VERTEX_SHADER, vertex_source_path)
        fragment_shader = Shader(gl.GL_FRAGMENT_SHADER, fragment_source_path)
        tcs_shader = Shader(gl.GL_TESS_CONTROL_SHADER, tcs_source_path)
        tes_shader = Shader(gl.GL_TESS_EVALUATION_SHADER, tes_source_path)
        geometry_shader = Shader(gl.GL_GEOMETRY_SHADER, geometry_source_path)
        super().__init__([vertex_shader, fragment_shader, tcs_shader, 
                          tes_shader, geometry_shader])
        vertex_shader.delete()
        fragment_shader.delete()
        tcs_shader.delete()
        tes_shader.delete()
        geometry_shader.delete()

        